"""
1. 定义模型要识别的类别
2. 定义模型要输入的数据
3. 创建模型
4. 获取模型输入数据
5. 将数据输入模型进行模型预测
6. 对预测结果进行非最大抑制
"""
import os
os.environ["TF_CPP_MIN_LOG_LEVEL"] = '2'
import warnings
warnings.filterwarnings("ignore")
import numpy as np
import matplotlib.pyplot as plt
from tensorflow.python.keras.preprocessing.image import load_img, img_to_array
from tensorflow.python.keras.applications.imagenet_utils import preprocess_input
from scipy.misc import imread

from nets.ssd_net import SSD300
from utils.ssd_utils import BBoxUtility


class SSDTest:
    def __init__(self):
        # 1. 定义模型要识别的类别
        self.classes_name = ['Aeroplane', 'Bicycle', 'Bird', 'Boat', 'Bottle',
                             'Bus', 'Car', 'Cat', 'Chair', 'Cow', 'Diningtable',
                             'Dog', 'Horse', 'Motorbike', 'Person', 'Pottedplant',
                             'Sheep', 'Sofa', 'Train', 'Tvmonitor']

        # 2. 定义模型要输入的数据
        self.classes_nums = len(self.classes_name) + 1
        self.input_shape = (300, 300, 3)

    def test(self):
        """
        模型测试
        :return:
        """
        # 创建模型
        model = SSD300(input_shape=self.input_shape, num_classes=self.classes_nums)

        # 加载模型
        model.load_weights("./ckpt/weights_SSD300.hdf5", by_name=True)

        # 获取数据
        feature = []
        images_binary = []
        for image in os.listdir("./images"):
            image_path = os.path.join("./images/", image)
            # 加载图片
            image = load_img(path=image_path, target_size=(self.input_shape[0], self.input_shape[1]))
            # 将图片转换为数组
            image = img_to_array(image)
            # 将读取的图片放入输入列表
            feature.append(image)

            # 将图片的二进制内容写入images，供后面图片可视化使用
            images_binary.append(imread(image_path))

        # print(np.asarray(feature).shape)  # (2, 300, 300, 3)

        # 获取模型的输入数据
        inputs = preprocess_input(np.asarray(feature))
        # print(inputs.shape)  # (2, 300, 300, 3)

        # 输入数据进行模型预测
        pre = model.predict(inputs)
        # (2, 7308, 33) 2：两张图片， 7308： 每张图片7308个default bboxes，
        # 33： 4(候选框的四个维度) + 21(21个预测的类别) + 8(default bboxes的参数)
        # print(pre.shape)

        # 对预测结果进行非最大抑制
        bb = BBoxUtility(self.classes_nums)
        res = bb.detection_out(pre)
        # 未加载模型的打印(200, 6) (200, 6): 每张图片经过NMS得到前200个default bbox，
        # 6：[label, confidence, xmin, ymin, xmax, ymax]
        # print(res[0].shape, res[1].shape)
        # # 加载训练好的模型打印：(132, 6) (26, 6)
        # print(res[0].shape, res[1].shape)

        return res, images_binary

    def filter_outputs(self, outputs, images):
        """
        对非最大化抑制结果进行过滤
        :param outputs: 经过SSD300得到的输出结果
        :param images: 图片的二进制数据
        :return:
        """
        for i, img in enumerate(images):
            # print("img:", type(img))  # mg: <class 'numpy.ndarray'>
            # print(outputs[:])
            # 将预测的label, confidence, xmin, ymin, xmax, ymax分离
            # 注意[:, i]和[:][i]的区别
            # [:, i]是取所有行的第i个维度，[:][i]是取[:]的第i个元素
            pre_label = outputs[i][:, 0]
            pre_conf = outputs[i][:, 1]
            pre_xmin = outputs[i][:, 2]
            pre_ymin = outputs[i][:, 3]
            pre_xmax = outputs[i][:, 4]
            pre_ymax = outputs[i][:, 5]
            # print(f"pre_label:{pre_label}, pre_conf:{pre_conf}, pre_xmin:{pre_xmin},"
            #       f" pre_ymin:{pre_ymin}, pre_xmax:{pre_xmax}, pre_ymax:{pre_ymax}")
            # print(type(pre_label))  # ndarray

            # 过滤指定置信度的候选框
            target_indices = [i for i, p in enumerate(pre_conf) if p >= 0.6]
            target_label = pre_label[target_indices]
            target_conf = pre_conf[target_indices]
            target_xmin = pre_xmin[target_indices]
            target_ymin = pre_ymin[target_indices]
            target_xmax = pre_xmax[target_indices]
            target_ymax = pre_ymax[target_indices]
            print(f"target_label:{target_label}, target_conf:{target_conf}, target_xmin:{target_xmin},"
                  f" target_ymin:{target_ymin}, target_xmax:{target_xmax}, target_ymax:{target_ymax}")

            # matplotlib可视化
            # 定义21中颜色，显示图片
            # currentAxis增加图中文本显示和标记显示
            colors = plt.cm.hsv(np.linspace(0, 1, 21)).tolist()  # shpae: (21, 4)
            plt.imshow(img / 255.)
            currentAxis = plt.gca()

            for i in range(target_conf.shape[0]):
                # round(89.3324, 开始四舍五入的位数), 返回第一个数按固定位置四舍五入的值
                xmin = int(round(target_xmin[i] * img.shape[1]))
                ymin = int(round(target_ymin[i] * img.shape[0]))
                xmax = int(round(target_xmax[i] * img.shape[1]))
                ymax = int(round(target_ymax[i] * img.shape[0]))

                # 获取该图片预测概率，名称，定义显示颜色
                score = target_conf[i]
                label = int(target_label[i])
                label_name = self.classes_name[label - 1]
                display_txt = '{:0.2f}, {}'.format(score, label_name)
                coords = (xmin, ymin), xmax - xmin + 1, ymax - ymin + 1
                color = colors[label]
                # 显示方框
                currentAxis.add_patch(plt.Rectangle(*coords, fill=False, edgecolor=color, linewidth=2))
                # 左上角显示概率以及名称
                currentAxis.text(xmin, ymin, display_txt, bbox={'facecolor': color, 'alpha': 0.5})

            plt.show()


if __name__ == '__main__':
    ssd = SSDTest()
    outputs, images = ssd.test()
    ssd.filter_outputs(outputs, images)
